from PySide2 import QtCore
from ImageConvert import *
from MVSDK import *
from PySide2.QtCore import *
import struct
import numpy as np

class Signals(QtCore.QObject):
    get_frame = QtCore.Signal(str)

Global_Signals = Signals()

g_cameraStatusUserInfo = b"statusInfo"

frameBuff_dict = {}
convertParams_dict = {}

class BITMAPFILEHEADER(Structure):
    _fields_ = [
                ('bfType', c_ushort),
                ('bfSize', c_uint),
                ('bfReserved1', c_ushort),
                ('bfReserved2', c_ushort),
                ('bfOffBits', c_uint),
                ]
 
class BITMAPINFOHEADER(Structure):
    _fields_ = [
                ('biSize', c_uint),
                ('biWidth', c_int),
                ('biHeight', c_int),
                ('biPlanes', c_ushort),
                ('biBitCount', c_ushort),
                ('biCompression', c_uint),
                ('biSizeImage', c_uint),
                ('biXPelsPerMeter', c_int),
                ('biYPelsPerMeter', c_int),
                ('biClrUsed', c_uint),
                ('biClrImportant', c_uint),
                ] 

# 调色板，只有8bit及以下才需要
# color palette
class RGBQUAD(Structure):
    _fields_ = [
                ('rgbBlue', c_ubyte),
                ('rgbGreen', c_ubyte),
                ('rgbRed', c_ubyte),
                ('rgbReserved', c_ubyte),
                ]

def Frame2cv(imageParams, userBuff):
    # 如果图像格式是 Mono8 直接使用
    # no format conversion required for Mono8
    if imageParams.pixelForamt == EPixelType.gvspPixelMono8:
        grayByteArray = bytearray(userBuff)
        return np.array(grayByteArray).reshape(imageParams.height, imageParams.width)
    else:
        # 转码 => BGR24
        # convert to BGR24
        rgbSize = c_int()
        rgbBuff = c_buffer(b'\0', imageParams.height * imageParams.width * 3)

        nRet = IMGCNV_ConvertToBGR24(cast(userBuff, c_void_p), \
                                     byref(imageParams), \
                                     cast(rgbBuff, c_void_p), \
                                     byref(rgbSize))

        colorByteArray = bytearray(rgbBuff)
        return np.array(colorByteArray).reshape(imageParams.height, imageParams.width, 3)

def WriteFrames(userInfo):
    global frameBuff_dict, convertParams_dict

    frame_id = 0
    for frameBuff, convertParams in zip(frameBuff_dict[str(c_char_p(userInfo).value)], convertParams_dict[str(c_char_p(userInfo).value)]):
        frame_id += 1
        # 保存bmp图片
        # save bmp image  
        bmpInfoHeader = BITMAPINFOHEADER() 
        bmpFileHeader = BITMAPFILEHEADER()
    
        uRgbQuadLen = 0
        rgbQuad = (RGBQUAD * 256)() # 调色板信息 | color palette info
        rgbBuff = c_buffer(b'\0', convertParams.height * convertParams.width * 3)
        
        # 如果图像格式是 Mono8 不需要转码
        # no format conversion required for Mono8
        if convertParams.pixelForamt == EPixelType.gvspPixelMono8:
            # 初始化调色板rgbQuad 实际应用中 rgbQuad 只需初始化一次
            # initialize color palette
            for i in range(0, 256):
                rgbQuad[i].rgbBlue = rgbQuad[i].rgbGreen = rgbQuad[i].rgbRed = i

            uRgbQuadLen = sizeof(RGBQUAD) * 256    
            bmpFileHeader.bfSize = sizeof(bmpFileHeader) + sizeof(bmpInfoHeader) + uRgbQuadLen + convertParams.dataSize
            bmpInfoHeader.biBitCount = 8
        else:
            # 转码 => BGR24
            # convert to BGR24
            rgbSize = c_int()
            nRet = IMGCNV_ConvertToBGR24(cast(frameBuff, c_void_p), byref(convertParams), \
                                        cast(rgbBuff, c_void_p), byref(rgbSize))
        
            if ( nRet != 0 ):
                print("image convert fail! errorCode = " + str(nRet))
                return -1 
            
            bmpFileHeader.bfSize = sizeof(bmpFileHeader) + sizeof(bmpInfoHeader) + rgbSize.value
            bmpInfoHeader.biBitCount = 24   
        
        bmpFileHeader.bfType = 0x4D42 # 文件头类型 'BM'(42 4D) | file type,  must be BM
        bmpFileHeader.bfReserved1 = 0 # 保留字 | reserved
        bmpFileHeader.bfReserved2 = 0 # 保留字 | reserved
        bmpFileHeader.bfOffBits = 54 + uRgbQuadLen # 位图像素数据的起始位置 | offset from the beginning of the BITMAPFILEHEADER structure to the bitmap bits
        
        bmpInfoHeader.biSize = 40     # 信息头所占字节数 | number of bytes required by the structure
        bmpInfoHeader.biWidth = convertParams.width
        bmpInfoHeader.biHeight = -convertParams.height
        bmpInfoHeader.biPlanes = 1    # 位图平面数| number of planes
        
        bmpInfoHeader.biCompression = 0 # 压缩类型，0 即不压缩 | compression type
        bmpInfoHeader.biSizeImage = 0
        bmpInfoHeader.biXPelsPerMeter = 0
        bmpInfoHeader.biYPelsPerMeter = 0
        bmpInfoHeader.biClrUsed = 0
        bmpInfoHeader.biClrImportant = 0    
        
        fileName = './image/%s-%d.bmp'%(str(c_char_p(userInfo).value), frame_id)
        imageFile = open(fileName, 'wb+') 
        
        imageFile.write(struct.pack('H', bmpFileHeader.bfType))
        imageFile.write(struct.pack('I', bmpFileHeader.bfSize))
        imageFile.write(struct.pack('H', bmpFileHeader.bfReserved1))
        imageFile.write(struct.pack('H', bmpFileHeader.bfReserved2))
        imageFile.write(struct.pack('I', bmpFileHeader.bfOffBits))
        
        imageFile.write(struct.pack('I', bmpInfoHeader.biSize))
        imageFile.write(struct.pack('i', bmpInfoHeader.biWidth))
        imageFile.write(struct.pack('i', bmpInfoHeader.biHeight))
        imageFile.write(struct.pack('H', bmpInfoHeader.biPlanes))
        imageFile.write(struct.pack('H', bmpInfoHeader.biBitCount))
        imageFile.write(struct.pack('I', bmpInfoHeader.biCompression))
        imageFile.write(struct.pack('I', bmpInfoHeader.biSizeImage))
        imageFile.write(struct.pack('i', bmpInfoHeader.biXPelsPerMeter))
        imageFile.write(struct.pack('i', bmpInfoHeader.biYPelsPerMeter))
        imageFile.write(struct.pack('I', bmpInfoHeader.biClrUsed))
        imageFile.write(struct.pack('I', bmpInfoHeader.biClrImportant))    
        
        if convertParams.pixelForamt == EPixelType.gvspPixelMono8:
            # 写入调色板信息
            # write out color palette info
            for i in range(0, 256):
                imageFile.write(struct.pack('B', rgbQuad[i].rgbBlue)) 
                imageFile.write(struct.pack('B', rgbQuad[i].rgbGreen))   
                imageFile.write(struct.pack('B', rgbQuad[i].rgbRed))           
                imageFile.write(struct.pack('B', rgbQuad[i].rgbReserved))
            
            imageFile.writelines(frameBuff)
        else:
            imageFile.writelines(rgbBuff)
            
        imageFile.close()
        print("save " + fileName + " success.")

def SaveFrame(frame, userInfo):
    global frameBuff_dict, convertParams_dict

    # 将裸数据图像拷出
    # copy image data out from frame
    imageSize = frame.contents.getImageSize(frame)
    buffAddr = frame.contents.getImage(frame)
    frameBuff = c_buffer(b'\0', imageSize)
    memmove(frameBuff, c_char_p(buffAddr), imageSize)
  
    # 给转码所需的参数赋值
    # fill conversion parameter
    convertParams = IMGCNV_SOpenParam()
    convertParams.dataSize = imageSize
    convertParams.height = frame.contents.getImageHeight(frame)
    convertParams.width = frame.contents.getImageWidth(frame)
    convertParams.paddingX = frame.contents.getImagePaddingX(frame)
    convertParams.paddingY = frame.contents.getImagePaddingY(frame)
    convertParams.pixelForamt = frame.contents.getImagePixelFormat(frame)
    
    if not str(c_char_p(userInfo).value) in convertParams_dict.keys():
        convertParams_dict[str(c_char_p(userInfo).value)] = []
    if not str(c_char_p(userInfo).value) in frameBuff_dict.keys():
        frameBuff_dict[str(c_char_p(userInfo).value)] = []
    convertParams_dict[str(c_char_p(userInfo).value)].append(convertParams)
    frameBuff_dict[str(c_char_p(userInfo).value)].append(frameBuff)

# 取流回调函数Ex
# grabbing callback function with userInfo parameter
def onGetFrameEx(frame, userInfo):
    
    nRet = frame.contents.valid(frame)
    if ( nRet != 0):
        print("frame is invalid!")
        # 释放驱动图像缓存资源
        # release frame resource before return
        frame.contents.release(frame)
        return         
 
    print("BlockId = %d userInfo = %s"  %(frame.contents.getBlockId(frame), c_char_p(userInfo).value))
    #此处客户应用程序应将图像拷贝出使用
    #Here you should copy image data out from frame for your own use
    
    SaveFrame(frame, userInfo)
    Global_Signals.get_frame.emit(True)

    # 释放驱动图像缓存资源
    # release frame resource at the end of use
    frame.contents.release(frame)

# 相机连接状态回调函数
# camera connection status change callback
def deviceLinkNotify(connectArg, linkInfo):
    if ( EVType.offLine == connectArg.contents.m_event ):
        print("camera has off line, userInfo [%s]" %(c_char_p(linkInfo).value))
    elif ( EVType.onLine == connectArg.contents.m_event ):
        print("camera has on line, userInfo [%s]" %(c_char_p(linkInfo).value))
    
connectCallBackFuncEx = connectCallBackEx(deviceLinkNotify)
frameCallbackFuncEx = callbackFuncEx(onGetFrameEx)

# 注册相机连接状态回调
# subscribe camera connection status change
def subscribeCameraStatus(camera):
    # 注册上下线通知
    # subscribe connection status notify
    eventSubscribe = pointer(GENICAM_EventSubscribe())
    eventSubscribeInfo = GENICAM_EventSubscribeInfo()
    eventSubscribeInfo.pCamera = pointer(camera)
    nRet = GENICAM_createEventSubscribe(byref(eventSubscribeInfo), byref(eventSubscribe))
    if ( nRet != 0):
        print("create eventSubscribe fail!")
        return -1
    
    nRet = eventSubscribe.contents.subscribeConnectArgsEx(eventSubscribe, connectCallBackFuncEx, g_cameraStatusUserInfo)
    if ( nRet != 0 ):
        print("subscribeConnectArgsEx fail!")
        # 释放相关资源
        # release subscribe resource before return
        eventSubscribe.contents.release(eventSubscribe)
        return -1  
    
    # 不再使用时，需释放相关资源
    # release subscribe resource at the end of use
    eventSubscribe.contents.release(eventSubscribe) 
    return 0
     
# 反注册相机连接状态回调
# unsubscribe camera connection status change
def unsubscribeCameraStatus(camera):
    # 反注册上下线通知
    # unsubscribe connection status notify
    eventSubscribe = pointer(GENICAM_EventSubscribe())
    eventSubscribeInfo = GENICAM_EventSubscribeInfo()
    eventSubscribeInfo.pCamera = pointer(camera)
    nRet = GENICAM_createEventSubscribe(byref(eventSubscribeInfo), byref(eventSubscribe))
    if ( nRet != 0):
        print("create eventSubscribe fail!")
        return -1
        
    nRet = eventSubscribe.contents.unsubscribeConnectArgsEx(eventSubscribe, connectCallBackFuncEx, g_cameraStatusUserInfo)
    if ( nRet != 0 ):
        print("unsubscribeConnectArgsEx fail!")
        # 释放相关资源
        # release subscribe resource before return
        eventSubscribe.contents.release(eventSubscribe)
        return -1
    
    # 不再使用时，需释放相关资源
    # release subscribe resource at the end of use
    eventSubscribe.contents.release(eventSubscribe)
    return 0   

# 打开相机
# open camera
def openCamera(camera):
    # 连接相机
    # connect camera
    nRet = camera.connect(camera, c_int(GENICAM_ECameraAccessPermission.accessPermissionControl))
    if ( nRet != 0 ):
        print("camera connect fail!")
        return -1
    else:
        print("camera connect success.")
  
    # 注册相机连接状态回调
    # subscribe camera connection status change
    nRet = subscribeCameraStatus(camera)
    if ( nRet != 0 ):
        print("subscribeCameraStatus fail!")
        return -1

    return 0

# 关闭相机
# close camera
def closeCamera(camera):
    # 反注册相机连接状态回调
    # unsubscribe camera connection status change
    nRet = unsubscribeCameraStatus(camera)
    if ( nRet != 0 ):
        print("unsubscribeCameraStatus fail!")
        return -1
  
    # 断开相机
    # disconnect camera
    nRet = camera.disConnect(byref(camera))
    if ( nRet != 0 ):
        print("disConnect camera fail!")
        return -1
    
    return 0    

# 设置曝光
# set camera ExposureTime
def setExposureTime(camera, dVal):
    # 通用属性设置:设置曝光 --根据属性类型，直接构造属性节点。如曝光是 double类型，构造doubleNode节点
    # create corresponding property node according to the value type of property, here is doubleNode
    exposureTimeNode = pointer(GENICAM_DoubleNode())
    exposureTimeNodeInfo = GENICAM_DoubleNodeInfo() 
    exposureTimeNodeInfo.pCamera = pointer(camera)
    exposureTimeNodeInfo.attrName = b"ExposureTime"
    nRet = GENICAM_createDoubleNode(byref(exposureTimeNodeInfo), byref(exposureTimeNode))
    if ( nRet != 0 ):
        print("create ExposureTime Node fail!")
        return -1
      
    # 设置曝光时间
    # set ExposureTime
    nRet = exposureTimeNode.contents.setValue(exposureTimeNode, c_double(dVal))  
    if ( nRet != 0 ):
        print("set ExposureTime value [%f]us fail!"  % (dVal))
        # 释放相关资源
        # release node resource before return
        exposureTimeNode.contents.release(exposureTimeNode)
        return -1
    else:
        print("set ExposureTime value [%f]us success." % (dVal))
            
    # 释放节点资源
    # release node resource at the end of use
    exposureTimeNode.contents.release(exposureTimeNode)    
    return 0
    
# 枚举相机
# enumerate camera
def enumCameras():
    # 获取系统单例
    # get system instance
    system = pointer(GENICAM_System())
    nRet = GENICAM_getSystemInstance(byref(system))
    if ( nRet != 0 ):
        print("getSystemInstance fail!")
        return None, None

    # 发现相机
    # discover camera 
    cameraList = pointer(GENICAM_Camera()) 
    cameraCnt = c_uint()
    nRet = system.contents.discovery(system, byref(cameraList), byref(cameraCnt), c_int(GENICAM_EProtocolType.typeAll))
    if ( nRet != 0 ):
        print("discovery fail!")
        return None, None
    elif cameraCnt.value < 1:
        print("discovery no camera!")
        return None, None
    else:
        print("cameraCnt: " + str(cameraCnt.value))
        return cameraCnt.value, cameraList
   